#include "MemoryManager.h"
#include <infiniband/verbs.h>
#include <iostream>

using std::cin;
using std::cout;
using std::endl;

MemoryPool::MemoryPool(void* addr, uint32_t size, uint32_t ck_num, uint32_t ck_size)
{
    base_ = addr;
    size_ = size;
    chunk_num_ = ck_num;
    chunk_size_ = ck_size;
    //init high performance stack
    chunk_allocator_.stack = new uint64_t[ck_num];
    chunk_allocator_.top = ck_num;
    uint64_t data_offset = (uint64_t)addr;
    for (uint32_t i = 0; i < chunk_num_; i++) {
        chunk_allocator_.stack[i] = data_offset;
        data_offset += chunk_size_;
    }
    return;
}

MemoryPool::~MemoryPool()
{
    if(chunk_allocator_.stack){
        delete [] chunk_allocator_.stack;
    }
}

void* MemoryPool::MemoryPoolAlloc() {
    if (chunk_allocator_.top == 0) {
        cout << "There is no chunk in memory pool" << "chunk size = " << chunk_size_ << endl;
        return NULL;
    }

    return (void*)chunk_allocator_.stack[--chunk_allocator_.top];
}

int MemoryPool::MemoryPoolFree(void *buffer) {
    if (((uint64_t)buffer >= (uint64_t)base_) && ((uint64_t)buffer <= ((uint64_t)base_ + size_))) 
    {
		if (chunk_allocator_.top == chunk_num_) 
        {
            cout << "Memory pool is Full" << ", chunk size = " << chunk_size_ << endl;
            return -1;
        }

        chunk_allocator_.stack[chunk_allocator_.top++] = (uint64_t)buffer;
		return 0;
	}
	else
	{
		return -1;
	}
    
}



MemoryManager::MemoryManager(void* memory, uint32_t memory_size, uint32_t max_msg_size)
: memory_(memory), memory_size_(memory_size), max_msg_size_(max_msg_size)
{ 
}

MemoryManager::~MemoryManager()
{
}

int MemoryManager::Init()
{
    uint32_t msg_total_size = (MSG_SIZE) * MSG_POOL_SIZE;
	uint32_t fragment_total_size = (FRAGMENT_SIZE) * FRAGMENT_POOL_SIZE;
    uint32_t ring_total_size = (max_msg_size_ + RING_SIZE) * RING_POOL_SIZE;

	if (memory_size_ < msg_total_size + fragment_total_size + ring_total_size) {
		cout << "memory size is too small, should more than " <<msg_total_size + fragment_total_size + ring_total_size << endl;
		return -1;
	}


    msg_pool_ = new MemoryPool(memory_, msg_total_size, MSG_POOL_SIZE, MSG_SIZE);
	if(!msg_pool_) 
    {
        cout << "Failed to create msg pool..." << endl;
        return -1;
    }
    fragment_pool_ = new MemoryPool((void*)((uint64_t)memory_ + msg_total_size), fragment_total_size, FRAGMENT_POOL_SIZE, FRAGMENT_SIZE);
    if(!fragment_pool_) 
    {
        cout << "Failed to create fragment pool..." << endl;
        return -1;
    }
    ring_pool_ = new MemoryPool((void*)((uint64_t)memory_ + msg_total_size + fragment_total_size), ring_total_size, RING_POOL_SIZE, max_msg_size_ + RING_SIZE);
    if(!ring_pool_) 
    {
        cout << "Failed to create ring pool..." << endl;
        return -1;
    }
    cout << "Succeed to init MemoryManager..."<<endl;
    return 0;
}

void MemoryManager::Destroy()
{
    ibv_dereg_mr(static_cast<ibv_mr*>(mr_));
    if(msg_pool_){
	    delete msg_pool_;
    }
    if(fragment_pool_){
        delete fragment_pool_;
    }
    if(ring_pool_){
        delete ring_pool_;
    }
}

void* MemoryManager::AllocFrag()
{
    return fragment_pool_->MemoryPoolAlloc();
}

void* MemoryManager::AllocMsg()
{
    return msg_pool_->MemoryPoolAlloc();
}

void* MemoryManager::AllocRing()
{
    return ring_pool_->MemoryPoolAlloc();
}

int MemoryManager::FreeFrag(void* buf)
{
    return fragment_pool_->MemoryPoolFree(buf);
}

int MemoryManager::FreeMsg(void* buf)
{
    return msg_pool_->MemoryPoolFree(buf);
}

int MemoryManager::FreeRing(void* buf)
{
    return ring_pool_->MemoryPoolFree(buf);
}
